from myutils.helper import get_item_append,get_accuracy_two,tqdm_update_epoch_batch_info

import gc
import torch
from tqdm import tqdm

from torch_ema import ExponentialMovingAverage

def train_siamese(dataLoader,model,loss,optimizer,scheduler,ema,is_test:bool,print_margin_batch:int,epoch_info:str,device,
        list_loss,list_accuracy_Y,list_accuracy_location):
    """单网络训练一个epoch

    Args:
        dataLoader:数据加载器
        model:模型
        loss:损失函数
        optimizer:优化器
        is_test:是否在测试（孪生姊妹网络中不能测试）
        print_margin_batch:每多少batch输出1次，暂不使用
        epoch_info:epoch信息
        
        list_history_one:单网络相关历史,[list_loss_one,list_loss_Y,list_loss_location,
                list_accuracy_Y,list_accuracy_location]

        list_history_siamese:孪生姊妹相关历史,[list_loss_siamese,list_loss_contrastive]
    """

    assert(not is_test)


    gc.collect()
    torch.cuda.empty_cache()
    
    
    ema = ExponentialMovingAverage(model.parameters(), decay=0.995)

    tqdm_dataLoader = tqdm(dataLoader,desc=epoch_info)
    # 批次,数据集 
    for batch_i, ( (img,shape,Y,location),(img_2,shape_2,Y_2,location_2) ) in enumerate(tqdm_dataLoader,1):
        
        # cuda: 垃圾收集，清理缓存，device转cuda
        gc.collect()
        torch.cuda.empty_cache()
        ((img  ,shape  ,Y  ,location)
        ,(img_2,shape_2,Y_2,location_2),) = (
            (img  .to(device),shape  .to(device),Y.  to(device),location  .to(device)),
            (img_2.to(device),shape_2.to(device),Y_2.to(device),location_2.to(device)), )
        
        # 正向传播
        (embedding,location_hat,Y_hat),(embedding_2,location_hat_2,Y_hat_2) = model(img,shape,img_2,shape_2)
        # 计算loss
        l = loss(
                embedding  ,Y_hat  ,Y  ,location_hat  ,location  , 
                embedding_2,Y_hat_2,Y_2,location_hat_2,location_2,)
        
        optimizer.zero_grad() # 梯度置为0
        l.backward()   # 反向传播
        optimizer.step()      # 梯度下降
        if scheduler: scheduler.step()
        if ema: ema.update()
        
        accuracy_Y        = get_accuracy_two(Y_hat,Y,Y_hat_2,Y_2)
        accuracy_location = get_accuracy_two(location_hat,location,location_hat_2,location_2)

        # 转原始类型(one):loss,accuracy
        get_item_append(list_loss,l)
        get_item_append(list_accuracy_Y,accuracy_Y)
        get_item_append(list_accuracy_location,accuracy_location)

        if batch_i%print_margin_batch==0:
            batch_info=f'{batch_i}/{len(dataLoader)}'
            tqdm_update_epoch_batch_info(tqdm_dataLoader,epoch_info,batch_info,print_margin_batch,
                    list_loss,list_accuracy_Y,list_accuracy_location)
            
    
    return